from datetime import timedelta
from zoneinfo import ZoneInfo

from flask import Flask, request, make_response
from cucc import *
from os.path import join, dirname, realpath, abspath, pardir
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers, RSAPrivateKey
from gmutil import SM2PrivateKey, SM2PublicKey, parse_pem_file
import json
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

SM4_SECRET_KEY = os.urandom(16)

with open(join(dirname(realpath(__file__)), 'THU_sm2_pri.pem'), 'rb') as f:
    THU_SM2_PRIVATE_KEY = SM2PrivateKey.from_pem(f.read())

with open(join(dirname(realpath(__file__)), 'THU_rsa_pri.pem'), 'rb') as f:
    THU_RSA_PRIVATE_KEY: RSAPrivateKey = serialization.load_pem_private_key(f.read(), None)

with open(join(dirname(realpath(__file__)), 'directory.json'), 'r', encoding='utf-8') as f:
    DIRECTORY = json.load(f)

UNIVERSITY_ID_THU = '4111010003'

app = Flask(__name__)


@app.route('/directory.json', methods=['GET'])
def directory():
    resp = make_response(json.dumps(DIRECTORY, ensure_ascii=False))
    resp.headers['Content-Type'] = 'application/json; charset=gb18030'
    return resp


@app.route('/qrcode/<userid>', methods=['GET'])
def qrcode(userid: str):
    return "CUCC:" + cucc_encode(bytes.fromhex(UNIVERSITY_ID_THU), userid.encode(), ALGORITHM_SM4_SM3, SM4_SECRET_KEY)


def id_num_hash(id_num:str, alg: str):
    salt = os.urandom(4)
    if alg == 'SM3':
        return 'SM3:' + base64.urlsafe_b64encode(sm3_hash(id_num.encode() + salt) + salt).decode()
    elif alg == 'SHA256':
        h = Hash(SHA256())
        h.update(id_num.encode())
        h.update(salt)
        return 'SHA256:' + base64.urlsafe_b64encode(h.finalize() + salt).decode()
    else:
        raise ValueError('Unknown hash algorithm')


def load_user_info(user_id):
    with open(join(dirname(realpath(__file__)), 'sample.json'), 'r', encoding='gb18030') as f:
        return json.load(f)


ERROR_REQUEST_PARAMETER = (0x01, '请求参数错误')
ERROR_ACQUIRER_ID = (0x02, '客户端高校代码未找到')
ERROR_ACQUIRER_KEY = (0x03, '客户端高校密钥未找到')
ERROR_REQUEST_ALGORITHM = (0x04, '客户端公钥算法类型不支持')
ERROR_REQUEST_SIGNATURE = (0x05, '请求内容签名验证错误')
ERROR_RESPONSE_ALGORITHM = (0x06, '请求的响应签名算法不支持')
ERROR_REQUEST_TIMESTAMP = (0x07, '请求时间戳错误')

ERROR_QRCODE_INVALID = (0x11, 'CUCC二维码无效')
ERROR_QRCODE_ISSUER = (0x12, 'CUCC二维码非本校发出')
ERROR_QRCODE_EXPIRED = (0x13, 'CUCC二维码过期')
ERROR_USER_INFO_NOT_FOUND = (0x21, 'CUCC二维码用户不存在')


@app.route('/auth', methods=['POST'])
def auth():
    def error_response(err, *args):
        logger.error('{}: {}'.format(err, '{} ' * len(args)).format(*args))
        ret = {k: err[i] for i, k in enumerate(('code', 'message'))}
        resp = make_response(ret)
        resp.headers['Content-Type'] = 'application/json; charset=gb18030'
        return resp

    try:
        acquirer_code = request.form['acquirer_code']
        barcode = request.form['barcode']
        timestamp = request.form['timestamp']
        kid = request.form['kid']
        signature = base64.urlsafe_b64decode(request.form['signature'] + '==')
        algorithm = request.form['algorithm'] if 'algorithm' in request.form else '01'
    except KeyError:
        return error_response(ERROR_REQUEST_PARAMETER)

    tz_cst = ZoneInfo('Asia/Shanghai')
    current_time = datetime.now(tz_cst)
    try:
        request_time = datetime.strptime(timestamp, '%Y%m%d%H%M%S').replace(tzinfo=tz_cst)
        offset = (current_time - request_time).total_seconds()
        if not -180 < offset < 180:
            return error_response(ERROR_REQUEST_TIMESTAMP, timestamp, offset)
    except ValueError:
        return error_response(ERROR_REQUEST_TIMESTAMP, timestamp)


    entry = next((e for e in DIRECTORY if e['code'] == acquirer_code), None)
    if entry is None:
        return error_response(ERROR_ACQUIRER_ID, acquirer_code)

    key = next((k for k in entry['public_keys']['keys'] if k['kid'] == kid), None)
    if key is None:
        return error_response(ERROR_ACQUIRER_KEY, acquirer_code, kid)

    to_sign = acquirer_code.encode() + barcode.encode() + timestamp.encode()
    if key['kty'] == 'EC' and key['crv'] == 'sm2p256v1' and key['alg'] == 'SGD_SM3_SM2':
        x_octets, y_octets = base64.urlsafe_b64decode(key['x']), base64.urlsafe_b64decode(key['y'])
        pubkey = SM2PublicKey.from_coordinates(
            int.from_bytes(x_octets, byteorder='big', signed=False),
            int.from_bytes(y_octets, byteorder='big', signed=False))
        if not pubkey.verify(to_sign, signature):
            return error_response(ERROR_REQUEST_SIGNATURE, 'SGD_SM3_SM2', acquirer_code, kid)
    elif key['kty'] == 'RSA' and key['alg'] == 'RS256':
        n = int.from_bytes(base64.urlsafe_b64decode(key['n']), byteorder='big', signed=False)
        e = int.from_bytes(base64.urlsafe_b64decode(key['e']), byteorder='big', signed=False)
        pubkey = RSAPublicNumbers(e, n).public_key()
        if not pubkey.verify(signature, to_sign, padding.PKCS1v15(), hashes.SHA256()):
            return error_response(ERROR_REQUEST_SIGNATURE, 'RS256', acquirer_code, kid)
    else:
        return error_response(ERROR_REQUEST_ALGORITHM, algorithm)

    try:
        university_id, timestamp, user_id = cucc_decode(barcode, SM4_SECRET_KEY)
    except ValueError as ve:
        return error_response(ERROR_QRCODE_INVALID)

    if university_id != UNIVERSITY_ID_THU:
        return error_response(ERROR_QRCODE_ISSUER)
    issued_time = datetime.strptime(timestamp, '%Y%m%d%H%M%S').replace(tzinfo=tz_cst)
    if issued_time > current_time or current_time - issued_time > timedelta(minutes=1):
        return error_response(ERROR_QRCODE_EXPIRED)

    user_info = load_user_info(user_id)
    # return error_response(ERROR_USER_INFO_NOT_FOUND)

    result = {
        'issuer_code': UNIVERSITY_ID_THU,
        'code': 0,
        'message': '认证成功',
        'timestamp': current_time.strftime('%Y%m%d%H%M%S'),
        'app_serial': base64.urlsafe_b64encode(bytes.fromhex(user_info['app_serial'])).decode(),
        'name': user_info['name'],
        'name_en': user_info['name_en'],
        'campus_id_number': user_info['campus_id_number'],
        'category': user_info['category'],
        'valid_since': user_info['valid_since'],
        'valid_until': user_info['valid_until'],
        'personal_id_hash': id_num_hash(user_info['id'], 'SM3'),
        'payment_token': user_info['payment_token']
    }

    user_info_items = ['{}={}'.format(k, v) for k, v in sorted(result.items(), key=lambda item: item[0])]

    to_sign = '&'.join(user_info_items)
    logger.debug("响应待签名内容：{}".format(to_sign))

    if algorithm == 'SM2':
        result['kid'] = '01'
        sig = THU_SM2_PRIVATE_KEY.sign(to_sign.encode('gb18030'))
        logger.debug("响应内容SM2签名值：{}".format(sig.hex()))
        result['signature'] = base64.urlsafe_b64encode(sig).decode().rstrip('=')
    elif algorithm == 'RSA':
        result['kid'] = '02'
        sig = THU_RSA_PRIVATE_KEY.sign(to_sign.encode('gb18030'), padding.PKCS1v15(), hashes.SHA256())
        logger.debug("响应内容RSA签名值：{}".format(sig.hex()))
        result['signature'] = base64.urlsafe_b64encode(sig).decode().rstrip('=')
    else:
        logger.error('请求签名的签名算法不支持：{}'.format(algorithm))
        return error_response(ERROR_RESPONSE_ALGORITHM, algorithm)

    response = make_response(json.dumps(result))
    response.headers['Content-Type'] = 'application/json; charset=gb18030'
    return response




if __name__ == "__main__":
    app.run(debug=False, host='0.0.0.0')
